from typing import Dict, Any, List

from autogen_core import CancellationToken

from mcp_tools.McpToolManager import get_tool_manager


class SqlTools:
    """基础SQL工具集合 - 提供通用SQL操作，业务SQL由LLM生成"""

    def __init__(self):
        self.tool_manager = None
        self._initialized = False

    async def initialize(self):
        """初始化工具管理器"""
        if not self._initialized:
            self.tool_manager = await get_tool_manager()
            self._initialized = True

    async def get_all_tables(self, cancellation_token: CancellationToken) -> Dict[str, Any]:
        """获取所有表信息"""
        try:
            await self.initialize()
            table_info = await self.tool_manager.query_database_tables(cancellation_token)
            return {
                "status": "success",
                "data": table_info,
                "description": "获取数据库中所有表的结构信息"
            }
        except Exception as e:
            return {
                "status": "error",
                "error": str(e),
                "description": "获取表信息失败"
            }

    async def get_table_structure(self, table_name: str, cancellation_token: CancellationToken) -> Dict[str, Any]:
        """获取指定表的结构信息"""
        try:
            await self.initialize()
            sql = f"DESCRIBE `{table_name}`;"
            result = await self.tool_manager.execute_sql_query(sql, cancellation_token)
            return {
                "status": "success",
                "data": result,
                "sql": sql,
                "description": f"获取表 {table_name} 的结构信息"
            }
        except Exception as e:
            return {
                "status": "error",
                "error": str(e),
                "description": f"获取表 {table_name} 结构失败"
            }

    async def get_table_sample_data(self, table_name: str, cancellation_token: CancellationToken, limit: int = 10) -> \
            Dict[str, Any]:
        """获取指定表的样本数据"""
        try:
            await self.initialize()
            sql = f"SELECT * FROM `{table_name}` LIMIT {limit};"
            result = await self.tool_manager.execute_sql_query(sql, cancellation_token)
            return {
                "status": "success",
                "data": result,
                "sql": sql,
                "description": f"获取表 {table_name} 的样本数据"
            }
        except Exception as e:
            return {
                "status": "error",
                "error": str(e),
                "description": f"获取表 {table_name} 样本数据失败"
            }

    async def execute_custom_sql(self, sql: str, cancellation_token: CancellationToken) -> Dict[str, Any]:
        """执行自定义SQL查询（由LLM生成）"""
        try:
            await self.initialize()
            result = await self.tool_manager.execute_sql_query(sql, cancellation_token)
            return {
                "status": "success",
                "data": result,
                "sql": sql,
                "description": "执行自定义SQL查询"
            }
        except Exception as e:
            return {
                "status": "error",
                "error": str(e),
                "description": "执行自定义SQL失败"
            }

    def get_available_tools(self) -> List[Dict[str, Any]]:
        """获取可用的工具列表"""
        return [
            {
                "name": "get_all_tables",
                "description": "获取数据库中所有表的结构信息",
                "parameters": {}
            },
            {
                "name": "get_table_structure",
                "description": "获取指定表的详细结构信息",
                "parameters": {
                    "table_name": "表名"
                }
            },
            {
                "name": "get_table_sample_data",
                "description": "获取指定表的样本数据",
                "parameters": {
                    "table_name": "表名",
                    "limit": "限制条数（可选，默认10）"
                }
            },
            {
                "name": "execute_custom_sql",
                "description": "执行自定义SQL查询（由LLM生成）",
                "parameters": {
                    "sql": "SQL查询语句"
                }
            }
        ]


# 全局SQL工具实例
sql_tools = SqlTools()


async def get_sql_tools() -> SqlTools:
    """获取全局SQL工具实例"""
    if not sql_tools._initialized:
        await sql_tools.initialize()
    return sql_tools
